import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
import chardet

# 1. 数据准备
# 检测文件编码
with open('../../附件2.csv', 'rb') as f:
    result = chardet.detect(f.read())
    encoding = result['encoding']

# 读取数据
data = pd.read_csv('../../附件2.csv', encoding=encoding)  # 使用检测到的编码
data['日期'] = pd.to_datetime(data['日期'], format='%Y/%m/%d')  # 转换日期格式
data = data.sort_values('日期')

# 2. 获取所有品类
categories = data['品类'].unique()

# 创建一个 DataFrame 用于存储所有品类的预测结果
all_forecasts = pd.DataFrame()

for category in categories:
    # 按品类筛选数据
    category_data = data[data['品类'] == category]
    category_data.set_index('日期', inplace=True)

    # 按日汇总销量
    daily_sales = category_data['销量'].resample('D').sum().reset_index()

    # 特征工程：创建时间特征
    daily_sales['day'] = daily_sales['日期'].dt.day
    daily_sales['month'] = daily_sales['日期'].dt.month
    daily_sales['year'] = daily_sales['日期'].dt.year
    daily_sales['dayofweek'] = daily_sales['日期'].dt.dayofweek

    # 划分特征和目标
    X = daily_sales.drop(columns=['日期', '销量'])
    y = daily_sales['销量']

    # 检查是否有足够的数据进行训练
    if len(X) < 2:  # 至少需要2个样本
        print(f"品类 {category} 的数据不足，跳过该品类。")
        continue

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)

    # 3. 随机森林模型训练
    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)

    # 4. 进行预测
    future_dates = pd.date_range(start='2024-07-01', periods=92, freq='D')
    future_features = pd.DataFrame({
        'day': future_dates.day,
        'month': future_dates.month,
        'year': future_dates.year,
        'dayofweek': future_dates.dayofweek
    })

    forecast = model.predict(future_features)
    forecast_series = pd.Series(forecast, index=future_dates)

    # 将预测结果添加到 DataFrame
    all_forecasts[category] = forecast_series

# 转置 DataFrame，使得行代表日期，列代表品类
all_forecasts = all_forecasts.transpose()
all_forecasts.reset_index(inplace=True)  # 重置索引
all_forecasts.columns = ['品类'] + list(all_forecasts.columns[1:])  # 设置列名

# 5. 保存所有品类的预测结果
with pd.ExcelWriter('RF销量预测结果.xlsx', engine='openpyxl', mode='w') as writer:
    all_forecasts.to_excel(writer, sheet_name='日销量预测结果', index=False)

print("所有品类的销量预测结果已成功保存到'RF销量预测结果.xlsx'文件中。")
